namespace Framework
{
    public class Repeater<T> : Decorator<T>
    {
        private int _repeatCount;
        private bool _endOnFailure;
        private int _iterationCount;

        public Repeater(int repeatCount, bool endOnFailure = false)
        {
            _repeatCount = repeatCount;
            _endOnFailure = endOnFailure;
        }

        public override void OnStart(T context)
        {
            base.OnStart(context);
            _iterationCount = 0;
        }

        public override TaskStatus Update(T context)
        {
            if (_repeatCount > 0 && _iterationCount == _repeatCount)
            {
                return TaskStatus.Success;
            }
            TaskStatus status = child.Tick(context);
            ++_iterationCount;
            if (_endOnFailure && status == TaskStatus.Failure)
            {
                return TaskStatus.Success;
            }
            return TaskStatus.Running;
        }
    }
}